/*
** Copyright (C) 2010 Calin Radoni
**
** This file is part of CHScanner.
**
** CHScanner is free software: you can redistribute it and/or modify
** it under the terms of the GNU General Public License as published by
** the Free Software Foundation, either version 3 of the License, or
** (at your option) any later version.
**
** CHScanner is distributed in the hope that it will be useful,
** but WITHOUT ANY WARRANTY; without even the implied warranty of
** MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
** GNU General Public License for more details.
**
** You should have received a copy of the GNU General Public License
** along with CHScanner. If not, see <http://www.gnu.org/licenses/>.
*/

#include "StdAfx.h"
#include "ScanMethodShutdown.h"

//--------------------------------------------------------------------

ScanMethodShutdown::ScanMethodShutdown(void) : ScanMethod()
{
	myType = __wmtShutdown;
	CleanUpUserCredentials();
	interactiveRun = true;
}

ScanMethodShutdown::~ScanMethodShutdown(void)
{
	CleanUpUserCredentials();
}

//--------------------------------------------------------------------

void ScanMethodShutdown::CleanUpUserCredentials(void)
{
	SecureZeroMemory(userName, (CREDUI_MAX_USERNAME_LENGTH + 1) * sizeof(TCHAR));
	SecureZeroMemory(userPass, (CREDUI_MAX_PASSWORD_LENGTH + 1) * sizeof(TCHAR));
}

void ScanMethodShutdown::SetUserCredentials(TCHAR* userNameIn, TCHAR* userPassIn, bool automaticModeIn)
{
	CleanUpUserCredentials();

	interactiveRun = !automaticModeIn;

	memcpy(userName, userNameIn, (CREDUI_MAX_USERNAME_LENGTH + 1) * sizeof(TCHAR));
	memcpy(userPass, userPassIn, (CREDUI_MAX_PASSWORD_LENGTH + 1) * sizeof(TCHAR));
}

//--------------------------------------------------------------------

int ScanMethodShutdown::InternalInit(void)
{
	CString str, sss;
	int pos;

	if(workPhase!=__wptScan)
		return MSR_ERR_PHASE;

	FillRandomBuffer();
	
	xReboot  = false;
	xForced  = true;
	xTimeOut = 30;
	xMessage = _T("");

	str = scanParameters.Trim();
	if(str.GetLength()!=0){
		if(str.GetLength()<2)
			return MSR_ERR_PARAMETER;

		// shutdown / reboot
		sss = str.Left(1);
		if(sss.CompareNoCase(_T("T"))==0) xReboot = false;
		else{
			if(sss.CompareNoCase(_T("F"))==0) xReboot = true;
			else return MSR_ERR_PARAMETER;
		}

		// forced
		sss = str.Mid(1,1);
		if(sss.CompareNoCase(_T("T"))==0) xForced = true;
		else{
			if(sss.CompareNoCase(_T("F"))==0) xForced = false;
			else return MSR_ERR_PARAMETER;
		}

		str = str.Mid(2);
		if(str.GetLength()>0){
			pos = str.Find(_T(" "));
			switch(pos){
				case -1:
					// only time-out specified
					sss = str;
					str.Empty();
					break;

				case 0 :
					// no time-out specified
					sss.Empty();
					str = str.Mid(1);
					break;

				default:
					sss = str.Left(pos);
					str = str.Mid(pos + 1);
					// time-out and message specified
					break;
			}

			if(!sss.IsEmpty())
				xTimeOut = _tstoi(sss);

			xMessage = str;
		}
	}

	if(xMessage.GetLength()>3072)
		xMessage = xMessage.Left(3072);

	if(xTimeOut<0)
		xTimeOut = 0;
	if(xTimeOut>MAX_SHUTDOWN_TIMEOUT)
		xTimeOut = MAX_SHUTDOWN_TIMEOUT;

	xUserHandle   = NULL;
	xThreadHandle = NULL;
	xTokenHandle  = NULL;
	xImpersonated = false;

	return MSR_ERR_OK;
}

int ScanMethodShutdown::InternalExecute(void)
{
	MMachineInfo* mmi;

	if(addrList==NULL) return MSR_ERR_PARAMETER;

	packSent = 0;

	if(!ElevatePrivileges()){
		abortReason = MSR_ERR_PRIV_ELEVATION;
		return MSR_ERR_OK;
	}

	ScanPair sp;
	while(GetNextPair(&sp, __sptIpOnly) && !stopStopStop){
		if(sp.addr!=NULL){
			if(!sp.addr->addrIP.IsEmpty()){
				if(workPhase==__wptDiscovery && sp.addr->alive){
					// it's the discovery phase and a previous discovery method has found it already
					// so skip this one
				}
				else{
					if(config->xMinDelay>0){
						if(lastScanTime!=0){
							waitDone = false;
							while(!waitDone){
								CheckResponses();
								SleepEx(1, TRUE);
								scanTime = GetTickCount();
								scanTime = ComputeDeltaTime(lastScanTime, scanTime);
								if(scanTime>=config->xMinDelay)
									waitDone = true;
							}
						}
						lastScanTime = GetTickCount();
					}

					if(ShutItDown(sp.addr->addrIP.GetS())){
						// succeded
						if(workPhase==__wptDiscovery){
							sp.addr->alive = true;
						}
						else{
							mmi = FindMachineByIp(sp.addr->addrIP.GetD());
							if(mmi!=NULL){
								mmi->address.alive = true;
								mmi->successfullyStopped = true;
							}
						}
					}

					packSent++;

					if((packSent % notifNetBios) == 0 && packSent>0)
						if(hWnd!=NULL)
							PostMessage(hWnd, CHS_MSG_PACKETS_SENT, 0, (LPARAM)packSent);
				}
			}
		}
	}

	RestorePrivileges();

	return MSR_ERR_OK;
}

//--------------------------------------------------------------------

bool ScanMethodShutdown::ElevatePrivileges(void)
{
	DWORD err;

	if(interactiveRun){
		if(!LogonUser(userName, _T(""), userPass, LOGON32_LOGON_NETWORK_CLEARTEXT, LOGON32_PROVIDER_DEFAULT, &xUserHandle)){
			return false;
		}
	
		if(!ImpersonateLoggedOnUser(xUserHandle)){
			CloseHandle(xUserHandle);
			xUserHandle = NULL;
			return false;
		}
		else xImpersonated = true;
	}
	else{
		xUserHandle = NULL;

		if(ImpersonateSelf(SecurityDelegation))
			xImpersonated = true;
	}

	xThreadHandle = GetCurrentThread();

	if(!OpenThreadToken(xThreadHandle, TOKEN_ALL_ACCESS /* TOKEN_ADJUST_PRIVILEGES | TOKEN_QUERY */, FALSE, &xTokenHandle)){
		RestorePrivileges();
		return false;
	}

	xPrivileges.Privileges[0].Luid.LowPart = 0;
	xPrivileges.Privileges[0].Luid.HighPart = 0;
	if(!LookupPrivilegeValue(NULL, SE_REMOTE_SHUTDOWN_NAME, &xPrivileges.Privileges[0].Luid)){
		RestorePrivileges();
		return false;
	}

	xPrivileges.PrivilegeCount = 1;
	xPrivileges.Privileges[0].Attributes = SE_PRIVILEGE_ENABLED;

	if(!AdjustTokenPrivileges(xTokenHandle, FALSE, &xPrivileges, 0, NULL, NULL)){
		RestorePrivileges();
		return false;
	}
	err = GetLastError(); // asa trebe, scrie in MSDN de ce
	switch(err){
		case ERROR_SUCCESS :
			// OK, go on
			break;

		case ERROR_NOT_ALL_ASSIGNED :
			RestorePrivileges();
			return false;
			break;

		default:
			RestorePrivileges();
			return false;
			break;
	}

	return true;
}

void ScanMethodShutdown::RestorePrivileges(void)
{
	if(xTokenHandle!=NULL){
		if(xPrivileges.Privileges[0].Attributes==SE_PRIVILEGE_ENABLED && xPrivileges.PrivilegeCount==1){
			xPrivileges.Privileges[0].Attributes = 0; // NU altceva, doar 0 face disable, scrie in MSDN !
			AdjustTokenPrivileges(xTokenHandle, FALSE, &xPrivileges, 0, NULL, NULL);
		}

		CloseHandle(xTokenHandle);
		xTokenHandle = NULL;
	}

	if(xThreadHandle!=NULL){
		CloseHandle(xThreadHandle);
		xThreadHandle = NULL;
	}

	if(xImpersonated){
		RevertToSelf();
		xImpersonated = false;
	}

	if(xUserHandle!=NULL){
		CloseHandle(xUserHandle);
		xUserHandle = NULL;
	}
}

bool ScanMethodShutdown::ShutItDown(CString ipStr)
{
	BOOL retVal;
	CT2A yIpStr(ipStr);
	CT2A yMessage(xMessage);

//	retVal = InitiateSystemShutdownEx(yIpStr, yMessage, xTimeOut, xForced, xReboot, SHTDN_REASON_MAJOR_SOFTWARE | SHTDN_REASON_MINOR_RECONFIG);
	retVal = InitiateSystemShutdown(
				yIpStr,
				yMessage,
				xTimeOut,
				xForced ? TRUE : FALSE,
				xReboot ? TRUE : FALSE
			);

	return retVal ? true : false;
}

//--------------------------------------------------------------------

void ScanMethodShutdown::ComputeRandomValues(void)
{
	// nothing to do
}

//--------------------------------------------------------------------
